import os
from typing import List
import threading
import tiktoken
from tqdm import trange
import time
import requests
import random
import json
# from langchain.document_loaders import PyPDFLoader
from langchain_community.document_loaders import PyPDFLoader

# 定义一个类 tokenCounter，用于计算文本的 token 数量和估算生成文本的价格
class tokenCounter():

    def __init__(self) -> None:
        # 初始化方法，设置编码方式和模型价格信息
        self.encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")  # 获取指定模型的编码方式
        self.model_price = {}  # 用于存储不同模型的价格信息
        
    # 计算字符串的 token 数量
    def num_tokens_from_string(self, string: str) -> int:
        return len(self.encoding.encode(string))  # 使用模型编码计算字符串的 token 数量

    # 计算字符串列表中所有字符串的总 token 数量
    def num_tokens_from_list_string(self, list_of_string: List[str]) -> int:
        num = 0
        for s in list_of_string:
            num += len(self.encoding.encode(s))  # 逐个计算每个字符串的 token 数量并累加
        return num
    
    # 计算生成文本的价格
    def compute_price(self, input_tokens, output_tokens, model):
        return (input_tokens / 1000) * self.model_price[model][0] + (output_tokens / 1000) * self.model_price[model][1]
        # 根据输入和输出的 token 数量及模型的价格计算生成文本的价格

    # 截断文本至指定最大长度的 token
    def text_truncation(self, text, max_len=1000):
        encoded_id = self.encoding.encode(text, disallowed_special=())  # 将文本进行编码
        return self.encoding.decode(encoded_id[:min(max_len, len(encoded_id))])  # 解码前 max_len 个 token，截取文本至指定长度

# 定义一个函数，用于加载 PDF 文件并返回截断后的文本

def load_pdf(file, max_len=1000):
    loader = PyPDFLoader(file)  # 使用 PyPDFLoader 加载 PDF 文件
    pages = loader.load_and_split()  # 加载并拆分 PDF 页面的内容
    encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")  # 获取指定模型的编码方式
    text = ''.join([p.page_content for p in pages])  # 将所有页面内容拼接为一个字符串
    return encoding.decode(encoding.encode(text)[:max_len])  # 将拼接后的内容编码后截断至指定长度，并解码返回

if __name__ == "__main__":
    # 初始化 tokenCounter 实例
    counter = tokenCounter()
    
    # 示例文本
    text = "这是一个用于测试的文本字符串。"
    num_tokens = counter.num_tokens_from_string(text)  # 计算文本的 token 数量
    print(f"文本的 token 数量: {num_tokens}")
    
    # 截断文本示例
    truncated_text = counter.text_truncation(text, max_len=5)  # 截断文本至最多 5 个 token
    print(f"截断后的文本: {truncated_text}")
    
    # 加载 PDF 文件示例
    pdf_text = load_pdf("./autosurvey.pdf", max_len=100)  # 加载并截断 PDF 文件内容
    print(f"PDF 文件的截断内容: {pdf_text}")